import sys
import random
from PyQt5.QtWidgets import QApplication, QMainWindow, QFrame, QMessageBox
from PyQt5.QtGui import QPainter, QColor, QBrush, QKeyEvent
from PyQt5.QtCore import QTimer, Qt

BLOCK_SIZE = 20
SCREEN_WIDTH = 640
SCREEN_HEIGHT = 480
SCREEN_W = SCREEN_WIDTH // BLOCK_SIZE
SCREEN_H = SCREEN_HEIGHT // BLOCK_SIZE
COLOR_BLACK = QColor(0, 0, 0)
COLOR_RED = QColor(255, 0, 0)
COLOR_GREEN = QColor(0, 255, 0)

class SnakeGame(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('贪吃蛇')
        # self.setGeometry(100, 100, SCREEN_WIDTH, SCREEN_HEIGHT)
        self.resize(SCREEN_WIDTH, SCREEN_HEIGHT)
        self.frame = GameFrame(self)
        self.setCentralWidget(self.frame)

class GameFrame(QFrame):
    def __init__(self, parent):
        super().__init__(parent)
        # 自动获取焦点
        self.setFocusPolicy(Qt.StrongFocus)
        self.initGame()
        
        # 创建定时器
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.updateGame)
        self.timer.start(200)

    def initGame(self):
        self.snake = [(10, 20), (10, 19), (10, 18)]
        self.direction = (1, 0)  # 向右
        self.food = self.placeFood()
        self.score = 0

    def placeFood(self):
        while True:
            # 随机生成食物
            food = (random.randint(0, 31), random.randint(0, 23))
            if food not in self.snake:
                # 食物不能在蛇身上
                return food

    def updateGame(self):
        head_x, head_y = self.snake[0]
        dir_x, dir_y = self.direction
        # 计算新的头部
        new_head = (head_x + dir_x, head_y + dir_y)

        # 检查碰撞
        if (new_head in self.snake or
            new_head[0] < 0 or new_head[0] >= 32 or
            new_head[1] < 0 or new_head[1] >= 24):
            QMessageBox.information(self, '游戏结束', f'得分: {self.score}')
            self.initGame()
            return

        # 更新蛇
        self.snake.insert(0, new_head)

        # 吃到食物
        if new_head == self.food:
            # 得分加1
            self.score += 1
            self.food = self.placeFood()
        else:
            # 没有吃到食物，删除尾部
            self.snake.pop()

        # 重绘
        self.update()

    def keyPressEvent(self, event: QKeyEvent):
        if event.key() == Qt.Key_Up and self.direction != (0, 1):
            self.direction = (0, -1)
        elif event.key() == Qt.Key_Down and self.direction != (0, -1):
            self.direction = (0, 1)
        elif event.key() == Qt.Key_Left and self.direction != (1, 0):
            self.direction = (-1, 0)
        elif event.key() == Qt.Key_Right and self.direction != (-1, 0):
            self.direction = (1, 0)
        elif event.key() == Qt.Key_Escape:
            self.initGame()
            self.update()

    def paintEvent(self, event):
        painter = QPainter(self)
        # 绘制背景
        painter.setBrush(QBrush(COLOR_BLACK))
        painter.drawRect(self.rect())
        # 绘制食物
        painter.setBrush(QBrush(COLOR_RED))
        painter.drawRect(self.food[0] * BLOCK_SIZE, self.food[1] * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE)
        # 绘制蛇
        painter.setBrush(QBrush(COLOR_GREEN))
        for segment in self.snake:
            painter.drawRect(segment[0] * BLOCK_SIZE, segment[1] * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    game = SnakeGame()
    game.show()
    sys.exit(app.exec_())